package com.tdy.cdc.util;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.alibaba.fastjson.serializer.ValueFilter;
import com.tdy.cdc.function.MySQLSinkFunction;
import org.apache.doris.flink.cfg.DorisExecutionOptions;
import org.apache.doris.flink.cfg.DorisOptions;
import org.apache.doris.flink.cfg.DorisReadOptions;
import org.apache.doris.flink.sink.DorisSink;
import org.apache.doris.flink.sink.writer.SimpleStringSerializer;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.connector.base.DeliveryGuarantee;
import org.apache.flink.connector.kafka.sink.KafkaRecordSerializationSchema;
import org.apache.flink.connector.kafka.sink.KafkaSink;
import org.apache.flink.connector.kafka.sink.TopicSelector;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkFixedPartitioner;
import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
import org.apache.hadoop.yarn.webapp.example.MyApp;
import org.apache.kafka.clients.producer.ProducerConfig;

import java.nio.charset.StandardCharsets;
import java.util.Properties;
import java.util.UUID;

/**
 * @author NanHuang
 * @Date 2023/3/1
 */
public class FlinkSinkUtil {
    public static final ValueFilter FILTER = (obj, s, v) -> {
        if(v==null)
            return "";
        return v;
    };
    public static DorisSink<String> getDorisSink(String tableName) {
        DorisSink.Builder<String> builder = DorisSink.builder();
        DorisOptions.Builder dorisBuilder = DorisOptions.builder();
        dorisBuilder.setFenodes(ConfigUtil.getProperty("doris.host"))
                .setTableIdentifier(tableName)
                .setUsername("admin")
                .setPassword("");



        DorisExecutionOptions.Builder  executionBuilder = DorisExecutionOptions.builder();
        executionBuilder.setLabelPrefix("label" + UUID.randomUUID()); //streamload label prefix
        Properties props = new Properties();
        props.setProperty("format", "json");
        props.setProperty("read_json_by_line", "true");
//        executionBuilder.setBufferCount(100000);
        executionBuilder.setBufferSize(100000);
        executionBuilder.setCheckInterval(2000);
        executionBuilder.setMaxRetries(3);
        executionBuilder.setStreamLoadProp(props);

        builder.setDorisReadOptions(DorisReadOptions.builder().build())
                .setDorisExecutionOptions(executionBuilder.build())
                .setSerializer(new SimpleStringSerializer()) //serialize according to string
                .setDorisOptions(dorisBuilder.build());
        return builder.build();
    }

    public static KafkaSink<JSONObject> getKafkaSink(Boolean isKey,String transactionalIdPrefix) {
        Properties properties = new Properties();
        properties.setProperty(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG,10*60*1000 + "");
        if (isKey) {
            // 有主键
            return KafkaSink.<JSONObject>builder()
                    .setBootstrapServers(ConfigUtil.getProperty("kafka.bootstrap.servers.test"))
                    .setRecordSerializer(
                            KafkaRecordSerializationSchema.builder()
                                    .setTopicSelector(new TopicSelector<JSONObject>() {
                                        @Override
                                        public String apply(JSONObject obj) {
                                            // 设置topic
                                            return obj.getString("table_name");
                                        }
                                    })
                                    .setValueSerializationSchema(new SerializationSchema<JSONObject>() {
                                        @Override
                                        public byte[] serialize(JSONObject obj) {
                                            // 定义value(处理FastJson转为字符串自动过滤null值问题)
                                            return JSON.toJSONString(obj,FILTER, SerializerFeature.WriteMapNullValue).getBytes(StandardCharsets.UTF_8);
                                        }
                                    })
                                    .setPartitioner(new FlinkKafkaPartitioner<JSONObject>() {
                                        @Override
                                        public int partition(JSONObject obj, byte[] bytes, byte[] bytes1, String s, int[] ints) {
                                            String key = obj.getString("pk");
                                            if (key == null) {
                                                key = obj.getString("id");
                                            }
                                            // 自定义分区策略
                                            return Math.abs(key.hashCode()) % ints.length;
                                        }
                                    })
                                    .setKeySerializationSchema(new SerializationSchema<JSONObject>() {
                                        @Override
                                        public byte[] serialize(JSONObject obj) {
                                            String key = obj.getString("pk");
                                            if (key == null) {
                                                key = obj.getString("id");
                                            }
                                            // 定义key
                                            return key.getBytes(StandardCharsets.UTF_8);
                                        }
                                    })
                                    .build()
                    )
                    .setKafkaProducerConfig(properties)
                    .setDeliverGuarantee(DeliveryGuarantee.EXACTLY_ONCE)
                    .setTransactionalIdPrefix(transactionalIdPrefix)
                    .build();
        }else {
            // 没有主键
            return KafkaSink.<JSONObject>builder()
                    .setBootstrapServers(ConfigUtil.getProperty("kafka.bootstrap.servers.test"))
                    .setRecordSerializer(
                            KafkaRecordSerializationSchema.builder()
                                    .setTopicSelector(new TopicSelector<JSONObject>() {
                                        @Override
                                        public String apply(JSONObject obj) {
                                            // 设置topic
                                            return obj.getString("table_name");
                                        }
                                    })
                                    .setValueSerializationSchema(new SerializationSchema<JSONObject>() {
                                        @Override
                                        public byte[] serialize(JSONObject obj) {
                                            // 定义value
                                            return JSON.toJSONString(obj,FILTER, SerializerFeature.WriteMapNullValue).getBytes(StandardCharsets.UTF_8);
                                        }
                                    })
                                    .build()
                    )
                    .setKafkaProducerConfig(properties)
                    .setDeliverGuarantee(DeliveryGuarantee.EXACTLY_ONCE)
                    .setTransactionalIdPrefix(transactionalIdPrefix)
                    .build();
        }
    }

    public static KafkaSink<String> getKafkaSink(String topicName) {
        Properties properties = new Properties();
        properties.setProperty(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG,10*60*1000 + "");
        return KafkaSink.<String>builder()
                .setBootstrapServers(ConfigUtil.getProperty("kafka.bootstrap.servers"))
                .setRecordSerializer(
                        KafkaRecordSerializationSchema.builder()
                                .setTopic(topicName)
                                .setValueSerializationSchema(new SerializationSchema<String>() {
                                    @Override
                                    public byte[] serialize(String s) {
                                        return s.getBytes(StandardCharsets.UTF_8);
                                    }
                                })
                                .build()
                )
                .setKafkaProducerConfig(properties)
                .setDeliverGuarantee(DeliveryGuarantee.EXACTLY_ONCE)
                .setTransactionalIdPrefix(topicName)
                .build();
    }

    public static String getKafkaWith(String topic,String ... format){
        String ft = "json";
        if (format.length != 0){
            ft = format[0];
        }
        return "with (" +
                "   'connector'='kafka'," +
                "   'properties.bootstrap.servers'='" + ConfigUtil.getProperty("kafka.bootstrap.servers.test") + "'," +
                "   'topic'='" + topic + "'," +
                "   'format'='" + ft + "'" +
                ")";
    }

    public static String getUpsertKafkaWith(String topic,String ... format){
        String ft = "json";
        if (format.length != 0){
            ft = format[0];
        }
        return "with (" +
                "       'connector'='upsert-kafka'," +
                "       'properties.bootstrap.servers'='" + ConfigUtil.getProperty("kafka.bootstrap.servers.test") + "'," +
                "       'topic'='" + topic + "'," +
                "       'key.format'='" + ft + "'," +
                "       'value.format'='" + ft + "'" +
                ")";
    }

    public static SinkFunction<JSONObject> getMySQLSink() {
        return new MySQLSinkFunction();
    }
}
